import onnx
from onnx import helper
from onnx import AttributeProto, TensorProto, GraphProto
import pdb

class RoundClipFixWrapper(object):

    def __init__(self, opVer=19):
        self.__opVer = opVer

    def build_model(self):
        X = helper.make_tensor_value_info('X', TensorProto.FLOAT, [1, 3, 512, 512])
        
        # Create one output (ValueInfoProto)
        Y = helper.make_tensor_value_info('Y', TensorProto.FLOAT, [1, 3, 512, 512])
        
        attrs = {
            'fix_input_bit': 'f32',
            'fix_output_bit': 'us8',
            'fix_min': 0,
            'fix_max': 255,
            'fix_scale': 255.18709,
            'fix_q_scale': 7,
            'fix_bitwidth_scale': 16,
            'fix_bottom_scale': 1.0,
            'fix_top_scale': 0.0039186934,
            'fix_quant_bitwidth': 0
        }

        # Create a node (NodeProto) - This is based on Pad-11
        nodeDef = helper.make_node(
            'RoundClipFix',             # op_type
            ['X'],                      # inputs
            ['Y'],                      # outputs
            name='round_clip_fix',
            **attrs
        )
        
        # Create the graph (GraphProto)
        graphDef = helper.make_graph(
            [nodeDef],         # nodes
            'test-model',      # name
            [X],  # inputs
            [Y],               # outputs
        )

        # Create the model (ModelProto)
        return helper.make_model(graphDef
            , opset_imports=[helper.make_opsetid("", self.__opVer)]
            , producer_name='ts_quant')
